import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt


def half_sine(x, A, w, phi, c):
    return A * np.sin(w * x + phi) + c


def Fit(data: np.array):
    # 提取x和y数据
    x_data = data[:, 0]  # 第一列作为x坐标
    y_data = data[:, 1]  # 第二列作为y坐标

    # 初始猜测值（A: 振幅, w: 角频率, phi: 相位偏移, c: 垂直偏移）
    initial_guess = [np.max(y_data) - np.min(y_data), 2 * np.pi / (x_data[-1] - x_data[0]), 0, np.mean(y_data)]

    # 使用 curve_fit 进行拟合，并增加最大迭代次数
    params, params_covariance = curve_fit(half_sine, x_data, y_data, p0=initial_guess, maxfev=5000)

    # 提取拟合参数
    A, w, phi, c = params

    # 打印拟合结果
    print(f"拟合参数: 振幅(A)={A:.2f}, 角频率(w)={w:.2f}, 相位偏移(phi)={phi:.2f}, 垂直偏移(c)={c:.2f}")

    # 生成用于绘图的 x 轴数据
    x_fit = np.linspace(x_data.min(), x_data.max(), 1000)
    y_fit = half_sine(x_fit, A, w, phi, c)

    # 绘制原始数据点
    plt.scatter(x_data, y_data, label='Data Points')
    # 绘制拟合曲线
    plt.plot(x_fit, y_fit, 'r-', label=f'Fit: A={A:.2f}, w={w:.2f}, φ={phi:.2f}, c={c:.2f}')
    # 添加图例
    plt.legend()
    # 设置图表标题和标签
    plt.title('Half Sine Curve Fitting')
    plt.xlabel('X-axis')
    plt.ylabel('Y-axis')
    plt.gca().invert_yaxis()

    return np.column_stack((x_fit, y_fit))
    # 显示图形
    # plt.show()
